#include "operation.h" 

namespace mslite {

Operation::Operation(cl_context context) {
    this->context = context;
}

Operation::~Operation() {
    cl_int status;
    for (auto mem: memObjs)
        status = clReleaseMemObject(mem.second);
    status = clReleaseKernel(kernel);
    status = clReleaseProgram(program);
}

cl_int Operation::setInput(const std::vector<std::vector<float>>& inputs, cl_command_queue commandQueue) {
    cl_int status;
    assert(inputs.size() == this->inputShapes.size());
    this->inputs.clear();
    for (int i=0; i<inputs.size(); i++) {
        int inputSize = getTotalSize(inputShapes[i]);
        std::string inputName = inputNames[i];
        assert(inputs[i].size() == inputSize);

        if (inputName == "weight" && clKernelName == "conv_1x1_image2d") {
            std::vector<float> i4o4weight = ReorderOHWIToI4O4(inputs[i], inputShapes[i], 2);
            this->inputs.push_back(i4o4weight);
        }
        else if (inputName == "weight" && clKernelName == "conv_1x1_tflite") {
            std::vector<float> i4o4weight = ReorderOHWIToI4O4(inputs[i], inputShapes[i], 2);
            this->inputs.push_back(i4o4weight);
        }
        else {
            this->inputs.push_back(inputs[i]);
        }

        
        if (memTypes[inputName] == Buffer) {
            cl_float *inputPtr = static_cast<cl_float *>(clEnqueueMapBuffer(commandQueue, memObjs[inputName], CL_TRUE, CL_MAP_WRITE,
                                                                            0, inputSize*sizeof(float), 0, NULL, NULL, &status));
            memcpy(inputPtr, this->inputs[i].data(), inputSize*sizeof(float));
            status = clEnqueueUnmapMemObject(commandQueue, memObjs[inputName], inputPtr, 0, NULL, NULL);
        }
        else if (memTypes[inputName] == Image2d) {
            size_t width = 0, height = 0;
            if (inputShapes[i].size() == 4) {       // 4-D tensor
                width = inputShapes[i][2];
                height = inputShapes[i][1] * inputShapes[i][3] / 4;
            }   
            else if (inputShapes[i].size() == 1) {
                width = inputShapes[i][0] / 4;
                height = 1;
            }
            size_t origin[3] = {0, 0, 0};           // 0, 0, 0
            size_t region[3] = {width, height, 1};  // width_, height_, 1
            status = clEnqueueWriteImage(commandQueue, memObjs[inputName], CL_TRUE, origin, region, 0, 0, inputs[i].data(), 0, NULL, NULL);
        }
    }
    return status;
}

cl_int Operation::getOutut(cl_command_queue commandQueue, std::vector<std::vector<float>> &output) {
    cl_int status;
    output.clear();
    output.resize(outputNames.size());
    
    for (int i=0; i<outputNames.size(); i++) {
        int osize = getTotalSize(this->outputShapes[i]);
        output[i].resize(osize);
        if (memTypes[outputNames[i]] == Buffer) {
            status = clEnqueueReadBuffer(commandQueue, memObjs[outputNames[i]], CL_TRUE, 0, sizeof(float)*osize, output[i].data(), 0, NULL, NULL);
        }
        else if (memTypes[outputNames[i]] == Image2d) {
            size_t width = 0, height = 0;
            if (outputShapes[i].size() == 4) {   // 4-D tensor
                width = outputShapes[i][2];
                height = outputShapes[i][1] * outputShapes[i][3] / 4;
            }   
            else if (outputShapes[i].size() == 1) {
                width = outputShapes[i][0] / 4;
                height = 1;
            }
            size_t origin[3] = {0, 0, 0};  // 0, 0, 0
            size_t region[3] = {width, height, 1};  // width_, height_, 1
            status = clEnqueueReadImage(commandQueue, memObjs[outputNames[i]], CL_TRUE, origin, region, 0, 0, output[i].data(), 0, NULL, NULL);
            CHECK(status, "Read Image Error: ");
        }
    }
    return status;
}

std::vector<std::vector<int>> Operation::getInputShapes(void) {
    return this->inputShapes;
}

cl_int Operation::setParams(std::vector<size_t> local) {
    this->localThreads = local;
    std::cout << "(" << localThreads[0] << " " << localThreads[1] << " " << localThreads[2] << ")";
    return 0;
}

cl_int Operation::run(cl_command_queue commandQueue, cl_event &profEvent) {
    cl_int status;
    status = clEnqueueNDRangeKernel(commandQueue, kernel, 3, NULL, globalThreads.data(), localThreads.data(), 0, NULL, &profEvent);
    clWaitForEvents(1, &profEvent);
    status = clFinish(commandQueue);
    return status;
}

}


